import json
import logging

from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

import streamlit as st
import pandas as pd

import DBTools as db

# 日志级别设置
logging.basicConfig(level=logging.INFO, filename="log.txt")

def print_json(data):
    """
    打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
    否则，直接打印该值。
    """
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list))):
        for item in data:
            print_json(item)
    elif (isinstance(data, (dict))):
        print(json.dumps(
            data,
            indent=4,
            ensure_ascii=False
        ))
    else:
        print(data)

# 头像配置
ICON_AI = '💻'
ICON_USER = '🧑'

#  描述数据库表结构
database_schema_string = """
CREATE TABLE orders (
    id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空
    customer_id INT NOT NULL, -- 客户ID，不允许为空
    product_id STR NOT NULL, -- 产品ID，不允许为空
    price DECIMAL(10,2) NOT NULL, -- 价格，不允许为空
    status INT NOT NULL, -- 订单状态，整数类型，不允许为空。0代表待支付，1代表已支付，2代表已退款
    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, -- 创建时间，默认为当前时间
    pay_time TIMESTAMP -- 支付时间，可以为空
);
"""


_ = load_dotenv(find_dotenv())
client = OpenAI()

def dspMessageHistory(role, content):
    print(f'''dspMessageHistory Start: {role}:{content}''')
    print_json(content)
    input_str = content
    with st.chat_message(role, avatar=ICON_AI if role == 'assistant' else ICON_USER):
        if role == "user":
            st.write(content)
        else:
            try:
                json_data = json.loads(input_str)
                if "dspStyle" in json_data and "dataSet" in json_data:
                    dsp_style = json_data["dspStyle"]
                    data_set = json_data["dataSet"]
                    data_frame = pd.DataFrame(data_set)
                    num_columns = data_frame.shape[1]
                    if num_columns != 2:
                        st.write("当前Demo仅支持二维数据展示，以表格展示数据")
                        st.table(data_frame)
                    else:            
                        categories = data_frame.iloc[:, 0]  # 第一列作为x轴数据（类别）  
                        values = data_frame.iloc[:, 1]      # 第二列作为y轴数据（值）  
    
                        # 将类别和值转换为适合st.bar_chart()的字典格式  
                        chart_data = {category: value for category, value in zip(categories, values)}  
                        
                        if dsp_style == "折线图":
                            st.line_chart(chart_data)
                        elif dsp_style == "柱状图":
                            st.bar_chart(chart_data)
                        elif dsp_style == "面积图":
                            st.area_chart(chart_data)
                        elif dsp_style == "散点图":
                            st.scatter_chart(chart_data)
                        else:
                            st.write("Demo暂时不支持的图表类型，以表格展示数据")
                            st.table(data_frame)
                else:
                    st.json(json_data)
            except json.JSONDecodeError:
                st.write(input_str)

def append_and_show(role, content):
    """
    将消息添加到st的messages列表中，并显示出来
    :param role: 角色，暂时就两个：assistant和user
    :param content: 消息内容
    :return:
    """

    st.session_state.messages.append({"role": role, "content": content})    
    #st.chat_message(role, avatar=ICON_AI if role == 'assistant' else ICON_USER).write(content)    
    dspMessageHistory(role, content)
    
def get_completion(prompt, model="gpt-3.5-turbo"):
    messages = prompt
    response = client.chat.completions.create(
        model = model,
        messages = messages,
        temperature=0,
        tools=[{  # 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb
            "type": "function",
            "function": {
                "name": "ask_database",
                "description": "Use this function to answer user questions about business. \
                            Output should be a fully formed SQL query.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": f"""
                            SQL query extracting info to answer the user's question.
                            SQL should be written using this database schema:
                            {database_schema_string}
                            The query should be returned in plain text, not in JSON.
                            The query should only contain grammars supported by SQLite.
                            """,
                        }
                    },
                    "required": ["query"],
                }
            }
        }],
    )
    return  response

#主程序
#初如化数据库
db.initDB(database_schema_string)

if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "我是你的信息助手，请问你查询什么信息？"}]

# 初始化thread到Streamlit的session_state中
if "thread" not in st.session_state:
    thread = client.beta.threads.create()
    logging.info(f'创建了新的thread: {thread.id}')
    st.session_state["thread"] = thread

# 将st中的messages列表中的消息显示出来
for msg in st.session_state.messages:
    #st.chat_message(msg["role"], avatar=ICON_AI if msg["role"] == 'assistant' else ICON_USER).write(msg["content"])
    dspMessageHistory(msg["role"], msg["content"])

if prompt := st.chat_input():
    append_and_show("user", prompt)
    
    templet = '''
    你是一个全能助手，判断用户的提示，
    如果需要可以基于 order 表回答用户问题。
    最终结果以json格式返回，其中使用名为dspStyle的key表示图表类型(折线图/柱状图/散点图/面积图)，
    使用名为dataSet的key表示数据集，
    数据集中的列名与值均为数据本身应具有的类型。
    例如以柱状图展示各商品的销售额，生成的json形式如下：
    {
        "dspStyle": "柱状图",
        "dataSet": [
            {"商品名":"商品1","销售额":123},
            {"商品名":"商品2","销售额":233},
            {"商品名":"商品3","销售额":167},
        ]
    }
    '''
    
    messages = [
        {"role": "system", "content": f"{templet}"},
        {"role": "user", "content": prompt}
    ]
    print(f'\nget_completion Call:')
    print_json(messages)
    response = get_completion(messages)
    if not hasattr(response, "content"):
        response.content = ""
    messages.append(response.choices[0].message)
    print(f'\nget_completion Return: ')
    print_json(response)
    print(f'\nmessages: {messages}')

    if response.choices[0].message.tool_calls is not None:
        tool_call = response.choices[0].message.tool_calls[0]
        if tool_call.function.name == "ask_database":
            arguments = tool_call.function.arguments
            args = json.loads(arguments)
            print("====SQL====")
            print(args["query"])
            result = db.ask_database(args["query"])
            print("====DB Records====")
            print(result)

            messages.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": "ask_database",
                "content": str(result)
            })
            print(f'\nget_completion Call:')
            print_json(messages)
            response = get_completion(messages)
            print(f'\nget_completion Return:')
            print_json(response)
            if hasattr(response, 'error'):
                append_and_show("assistant", response.error['message'])
            elif response.choices[0].message.content is not None:
                append_and_show("assistant", response.choices[0].message.content)   
            else:
                print(response)
    elif response.choices[0].message.content is not None:
        append_and_show("assistant", response.choices[0].message.content)
    else:
        append_and_show("assistant", "对不起，我出错了哦。")
    
     # 创建OpenAI的message
    #message = client.beta.threads.messages.create(thread_id=st.session_state.thread.id, role="user", content=prompt)
